﻿#include "OrszagTangVortex.h"
#include "Algorithm.h"
#include "DataObject.h"
#include "Mesh.h"
#include <vtkm/VectorAnalysis.h>

#include <vtkm/worklet/WorkletMapField.h>
#include <vtkm/worklet/WorkletMapTopology.h>

#include <vtkm/Math.h>
#include <vtkm/cont/ArrayHandleGroupVec.h>
#include <vtkm/exec/ParametricCoordinates.h>

struct UpdateVelocityAndPositionOTVortex : vtkm::worklet::WorkletMapField
{
  using ControlSignature = void(FieldIn bnd,
                                FieldIn point_mass,
                                FieldInOut force,
                                FieldInOut velocity,
                                FieldInOut position);
  using ExecutionSignature = void(_1, _2, _3, _4, _5);

  UpdateVelocityAndPositionOTVortex(Real dt)
    : _dt(dt)
  {
  }
  template<typename CoordType>
  VTKM_EXEC void operator()(const Id& bid,
                            const Real& point_mass,
                            Vec3f& force,
                            Vec3f& vel,
                            CoordType& point) const
  {
    //if (bid == -1)
    //{
    //  //vel += force * _dt / point_mass;
    //  //point += vel * _dt;
    //}
    //else if (bid != 4)
    //{
    //  Vec3f normal{ 0, 0, 0 };
    //  if (bid == 0)
    //  {
    //    normal = { 0, -1, 0 };
    //  }
    //  else if (bid == 1)
    //  {
    //    normal = { 1, 0, 0 };
    //  }
    //  else if (bid == 2)
    //  {
    //    normal = { 0, 1, 0 };
    //  }
    //  else if (bid == 3)
    //  {
    //    normal = { -1, 0, 0 };
    //  }
    //  force = (force - vtkm::Dot(force, normal) * normal);
    //  //vel = vel - vtkm::Dot(vel, normal) * normal;
    //  //point += vel * _dt;
    //}
    //else
    //{
    //  force = Vec3f{0,0,0};
    //}
    if (bid == 0 || bid == 2)
    {
      //force[1] = 0;
      vel += force * _dt / point_mass;
      point += vel * _dt;
      //std::cout << "force: " << force << std::endl;
      //std::cout << "vel: " << vel << std::endl;
      //std::cout << "point: " << point << std::endl;
    }
    else if (bid == 1 || bid == 3)
    {
      //force[0] = 0;
      vel += force * _dt / point_mass;
      point += vel * _dt;
    }
    else if (bid == 4)
    {
      //force[0] = 0;
      //force[1] = 0;
      vel += force * _dt / point_mass;
      point += vel * _dt;
    }
    else if (bid == -1)
    {
      vel += force * _dt / point_mass;
      point += vel * _dt;
    }
    else
    {
      this->RaiseError("Undefined bid");
    }
  }

  Real _dt;
};

struct UpdateCellEnergyOTVortex : WorkletVisitCellsWithPoints
{
  // clang-format off
  using ControlSignature =  void(CellSetIn, FieldInCell mass, FieldInPoint velocity, FieldInCell force, FieldInOutCell energy,  FieldInPoint pts);
  // clang-format on
  using ExecutionSignature = void(PointCount, _2, _3, _4, _5, _6);

  UpdateCellEnergyOTVortex(Real dt)
    : _dt(dt)
  {
  }
  template<typename ForceType, typename VelocityType, typename PType>
  VTKM_EXEC void operator()(const IdComponent& num_pts,
                            const Real& mass,
                            const VelocityType& velocity,
                            const ForceType& force,
                            Real& energy,
                            const PType& pts) const
  {
    Real rhs = 0;
    for (IdComponent i = 0; i < num_pts; i++)
    {
      rhs += vtkm::Dot(force[i], velocity[i]);
    }
    auto p = (pts[0] + pts[1] + pts[2]) / 3;

    auto x = p[0];
    auto y = p[1];
    Real source = 0;
    energy += _dt * (-rhs / mass + source);
  }

  Real _dt;
};

struct InitialConditionOnCellOTVortex : vtkm::worklet::WorkletVisitCellsWithPoints
{
  using ValueType = Vec<FloatDefault, 5>;

  using ControlSignature = void(CellSetIn, FieldInPoint, FieldOut density, FieldOut pressure);
  using ExecutionSignature = void(CellShape, _2, _3, _4);

  template<typename CellShapeTag, typename PType>
  VTKM_EXEC void operator()(const CellShapeTag& shape,
                            const PType& pts,
                            FloatDefault& density,
                            FloatDefault& pressure) const
  {
    auto p = (pts[0] + pts[1] + pts[2]) / 3;

    auto x = p[0];
    auto y = p[1];
    constexpr auto pi = vtkm::Pi();
    density = 25 / (36 * pi);
    pressure = 5 / (12 * pi);
  }
};

struct InitialConditionOnPointOTVortex : vtkm::worklet::WorkletMapField
{
  using ControlSignature = void(FieldIn point, FieldOut velocity);
  using ExecutionSignature = void(_1, _2);

  VTKM_EXEC
  void operator()(const Vec3f& p, Vec3f& velocity) const
  {
    constexpr auto pi = vtkm::Pi();
    auto x = p[0];
    auto y = p[1];
    Real u = -vtkm::Sin(2 * pi * y);
    Real v = vtkm::Sin(2 * pi * x);
    velocity = Vec3f{ u, v, 0 };
  }
};

OrszagTangVortex::OrszagTangVortex(const Configuration& cfg)
  : SGHSystem(cfg)
{
}

RegisterObject(OrszagTangVortex);
void OrszagTangVortex::InitialCondition()
{
  auto& cellset = _data_set.GetCellSet();
  auto& coords = _data_set.GetCoordinateSystem();
  auto& density = GetVariableArray<Real>("density");
  auto& pressure = GetVariableArray<Real>("pressure");
  auto& vel = GetVariableArray<Vec3f>("velocity");

  Invoker{}(InitialConditionOnCellOTVortex{}, cellset, coords, density, pressure);
  Invoker{}(InitialConditionOnPointOTVortex{}, coords, vel);


}

void OrszagTangVortex::Solve()
{
  ComputeTimeStepper(); // 计算 dt

  auto& point_mass = GetVariableArray<Real>("point_mass");
  auto& force = GetVariableArray<Vec3f>("force");
  auto& velocity = GetVariableArray<Vec3f>("velocity");
  auto& position = _mesh.GetCoordinate();
  auto& bnd_id = _mesh.GetBoundaryId();

  invoker(UpdateVelocityAndPositionOTVortex{ _dt }, bnd_id, point_mass, force, velocity, position);

  auto cell_mass = GetVariableArray<Real>("cell_mass");
  auto energy = GetVariableArray<Real>("energy");

  invoker(UpdateCellEnergyOTVortex{ _dt },
          _mesh.GetCellSet(),
          cell_mass,
          velocity,
          vtkm::cont::make_ArrayHandleGroupVec<3>(_c2n_force),
          energy,
          _mesh.GetCoordinate());
}
